from typing import Iterable

from prometheus_client import Gauge, CollectorRegistry, Counter


class FlowGatewayMetrics:

    def __init__(self):
        self.registry = CollectorRegistry(auto_describe=True)

    def _create_gauge(self, name: str, documentation: str, labelnames: Iterable[str] = ()) -> Gauge:
        return Gauge(name=name, documentation=documentation, labelnames=labelnames, registry=self.registry)

    def _create_counter(self, name: str, documentation: str, labelnames: Iterable[str] = ()) -> Counter:
        return Counter(name=name, documentation=documentation, labelnames=labelnames, registry=self.registry)

    def export_received_data_bytes_total(self, labels: dict, network_status):
        received_data_bytes_total = self._create_counter("received_data_bytes_total",
                                                         "自服务启动以来已接收的数据量，单位是字节。",
                                                         list(labels.keys()))
        received_data_bytes_total.labels(**labels)._value.set(network_status.received_data_bytes_total)

    def export_sent_data_bytes_total(self, labels: dict, network_status):
        sent_data_bytes_total = self._create_counter("sent_data_bytes_total",
                                                     "自服务启动以来已发送的数据量，单位是字节。",
                                                     list(labels.keys()))
        sent_data_bytes_total.labels(**labels)._value.set(network_status.sent_data_bytes_total)

    def export_downlink_bandwidth_usage(self, labels: dict, network_status):
        downlink_bandwidth_usage = self._create_gauge("downlink_bandwidth_usage", "当前下行带宽使用情况，单位是字节。",
                                                      list(labels.keys()))
        downlink_bandwidth_usage.labels(**labels).set(network_status.downlink_bandwidth_usage)

    def export_uplink_bandwidth_usage(self, labels: dict, network_status):
        uplink_bandwidth_usage = self._create_gauge("uplink_bandwidth_usage", "当前上行带宽使用情况，单位是字节。",
                                                    list(labels.keys()))
        uplink_bandwidth_usage.labels(**labels).set(network_status.uplink_bandwidth_usage)

    def export_device_heartbeat(self, device_number: str, timestamp: int):
        device_heartbeat = self._create_counter("device_heartbeat", "设备心跳，每分钟一次", ['device_number'])
        device_heartbeat.labels(device_number=device_number)._value.set(timestamp)
